#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/9/15 16:33
# @Author  : MAI Haoyuan
# @Site    :
# @File    : multi_factor_process.py
# @Software: PyCharm

"""
多因子筛选并合成脚本
"""

import pandas as pd
import numpy as np
from quant_researcher.quant.datasource_fetch.factor_api import factor_exposure_related
from quant_researcher.quant.datasource_fetch.index_api import index_components
import quant_researcher.quant.factors.factor_preprocess.preprocess as fp
from quant_researcher.quant.project_tool import hammer
from sklearn.linear_model import LinearRegression
from sklearn import decomposition


def multi_get_factor_data(factor, stock_pool_list, period_start, period_end):
    # temp = factor_exposure_related.get_stock_factor_exposure(factor, stock_pool_list, period_start, period_end)
    temp = factor_exposure_related.get_stock_factor_exposure(factor, stock_pool_list, period_start, period_end)
    temp = temp.rename(columns={'end_date': 'tradedate', 'stock_code': 'stockcode'})
    temp = temp.set_index(['stockcode', 'tradedate'])[factor].unstack(level=0).ffill().stack().rename(
        factor).reset_index()
    temp['factor'] = factor
    temp = temp.rename(columns={factor: 'value'})
    return temp


def multi_factor_process(factor_list=None, factor_valid_num=5, period_start='2020-01-01',
                         period_end='2020-05-01', asset_pool='HS300', asset_type='stock',
                         filter_indicator=None, method='rank_sum', winsor=True, standard=True,
                         multi_fetch=True, **kwargs):
    """
    多因子合成函数

    :param list list factor_list: 股票因子列表
    :param int factor_valid_num: 默认为5个参与合成
    :param str asset_pool:
    :param str asset_type:
    :param str period_start: '2020-01-01'
    :param str period_end: '2020-01-01'
    :param str stock_pool: 支持'HS300', 'ALLA'
    :param filter_indicator: 因子筛选指标，包含各个日期不同因子的某一指标值，用于排序筛选。例如根据因子拥挤度筛选因子 # todo 没有传入筛选指标怎么办？
        tradedate   |roe   |pb
        2010-01-03 | 0.2  |4.2
        2010-02-01| 0.4  |4.2
    :param method: 多因子合成方法，现在支持rank_sum（排序相加）, rank_product（排序相乘）,
               external_weight（外部权重加权）, regression（回归预测收益率）,pca（回归pca分解）
    :param bool winsor : 因子是否去极值
    :param bool standard: 因子是否标准化
    :param multi_fetch:
    :param kwargs:
        1. factor_data: 外部因子数据
            一个包含ROE，PB指标的股票因子数据结构如下，其中标的代码由asset_type+code组成：
            —————————————
            stockcode|tradedate|roe  |pb
            000001|2010-01-03| 0.2  |5.2
            000002|2010-01-03| 0.15|4.2
        2. weight_indicator： 加权指标数据
            加权指标数据包含的因子需要和因子数据包含的因子一致。
            —————————————
            tradedate   |roe   |pb
            2010-01-03| 0.2  |4.2
            2010-02-01| 0.4  |4.2
        3. all_asset_quote：资产行情数据。
            资产行情数据需包含因子数据同一时间区间内的所有标的数据。
            stockcode|tradedate|close
            000001|2010-01-03| 14.33
            000002|2010-01-03| 25.50
        4. orthogonalize： 因子之间是否对称正交，只适用于regression和pca方法，包含所有因子。

    :return:  返回合成后的因子
    """

    method_list = ['rank_sum', 'rank_product', 'external_weight', 'regression', 'pca']
    universe_dict = {'SH50': '000016', 'SH180': '000010', 'HS300': '399300', 'ZZ500': '000905',
                     'ZZ800': '000906', 'CYBZ': '399006', 'ZXB': '399005', 'ALLA': '000985'}
    assert method in method_list, f"请输入支持的方法，目前支持的方法有:{', '.join(method_list)}"
    stock_pool_list = index_components.get_index_hist_stock_pool(universe_dict[asset_pool], period_start, period_end)['stock_code'].tolist()
    if factor_list is None:
        assert 'factor_data' in kwargs, 'factor_list为空，请传入factor_data变量'
        factor_data_df = kwargs['factor_data']
        factor_names = factor_data_df.columns.difference(['tradedate', f'{asset_type}code']).tolist()
    else:
        if multi_fetch:
            info = {
                'func': multi_get_factor_data,
                'stock_pool_list': stock_pool_list,
                'period_start': period_start,
                'period_end': period_end,
                # 'prepare_these_conn': ['tk_factors']
            }
            factor_data_list = hammer.split_a_list_into_batches_and_deal_each_with_a_process(
                the_list=factor_list, info=info, ignore_err=True, how_many=len(factor_list), what_version='v1',

            )
        else:
            factor_data_list = []
            for factor in factor_list:
                # factor_table_name = factor_info.get_factor_table_name(factor)[0]
                # temp = factor_exposure_related.get_stock_factor_exposure(factor, stock_pool_list, period_start,
                #                                                          period_end)
                temp = factor_exposure_related.get_stock_factor_exposure(factor, stock_pool_list, period_start,
                                                                         period_end)
                temp = temp.rename(columns={'end_date': 'tradedate', 'stock_code': 'stockcode'})
                temp = temp.set_index(['stockcode','tradedate'])[factor].unstack(level=0).ffill().stack().rename(factor).reset_index()
                temp['factor'] = factor
                temp = temp.rename(columns={factor: 'value'})
                factor_data_list.append(temp)
        factor_data_df = pd.concat(factor_data_list).set_index(['tradedate', 'stockcode', 'factor'])[
            'value'].unstack().reset_index()
        factor_data_df.tradedate = pd.to_datetime(factor_data_df.tradedate)
        factor_names = factor_list[:]
    filter_rank = filter_indicator.set_index('tradedate').rank(axis=1, ascending=False)
    filter_mask = filter_rank <= factor_valid_num
    factor_dates_index = pd.Index(factor_data_df.tradedate.unique()).union(pd.to_datetime(filter_mask.index))

    filter_mask = filter_mask.reindex(factor_dates_index).ffill()
    filter_mask.index = filter_mask.index.rename('tradedate')
    if asset_type == 'manager':
        fac_data = factor_data_df.set_index(['tradedate', f'{asset_type}code', 'manager_code', 'fund_type'])
    else:
        fac_data = factor_data_df.set_index(['tradedate', f'{asset_type}code'])
    if winsor:
        fac_data = fac_data.groupby('tradedate').apply(lambda x: x.apply(fp.mstats.winsorize))
    if standard:
        fac_data = fac_data.groupby('tradedate').apply(lambda x: x.apply(fp.standardize))
    if method == 'rank_sum':
        fac_data = fac_data.groupby('tradedate').rank(pct=True, axis=0, na_option='top')
        fac_data = fac_data.groupby(f'{asset_type}code').apply(lambda x: x.where(filter_mask))
        fac_data = fac_data.mean(axis=1).rename('合成因子')
    elif method == 'rank_product':
        fac_data = fac_data.groupby('tradedate').rank(pct=True, axis=0, na_option='top')
        fac_data = fac_data.groupby(f'{asset_type}code').apply(lambda x: x.where(filter_mask))
        fac_data = fac_data.product(axis=1).rename('合成因子')
    elif method == 'external_weight':
        assert 'weight_indicator' in kwargs, '采用加权方式合成因子，请传入weight_indicator变量'
        weight_df = kwargs['weight_indicator'].set_index('tradedate')
        weight_factor_names = weight_df.columns.tolist()
        assert set(weight_factor_names) == set(factor_names), "请确认输入的权重数据包含的因子与因子数据一致"
        weight_df = weight_df.shift().where(filter_mask).fillna(0).reindex(
            factor_dates_index).ffill().abs()  # 当期只有上一期的ic值,所以要shift一下
        weight_df = weight_df.div(weight_df.sum(axis=1), axis=0).loc[period_start:]  # nan填充0，认为当期该因子未被选中
        weight_df.index = weight_df.index.rename('tradedate')
        # if orthogonalize:  正交化方法不能简单的应用到每期不同因子的情况，对于包含某个因子的5个因子正交与10个因子正交的值可能不同。
        #     fac_data = fp.symmetric_orth(fac_data, asset_type)
        fac_data = fac_data.groupby(level=1).apply(lambda x: x * weight_df).sum(axis=1).rename('合成因子')
    elif method == 'regression':
        assert 'all_asset_quote' in kwargs, '请传入资产行情（all_asset_quote）变量，用于回归'
        regress_dates = factor_data_df.tradedate.unique()
        asset_codes = factor_data_df[f'{asset_type}code'].unique()
        stock_price_indexed = kwargs['all_asset_quote'].set_index(['tradedate', f'{asset_type}code'])
        stock_returns = stock_price_indexed.groupby(level=1).pct_change()
        # 填充NaN，同时做对称正交化
        fac_data_filled = fp.symmetric_orth(fac_data, asset_type)
        new_factor = pd.DataFrame()

        for i in range(len(regress_dates) - 1):
            df = pd.DataFrame()
            date = regress_dates[i]
            date_forward = regress_dates[i + 1]
            y = stock_returns.loc[(date_forward,),]
            X = fac_data_filled.loc[(date,),]
            if np.any(np.isnan(X)):
                df = pd.DataFrame()
                df[f'{asset_type}code'] = asset_codes
                df['合成因子'] = np.mean(new_factor['合成因子'])
                df['tradedate'] = date_forward
                new_factor = pd.concat([new_factor, df])
                continue
            model = LinearRegression().fit(X, y)
            df[f'{asset_type}code'] = asset_codes
            df['合成因子'] = model.predict(fac_data_filled.loc[(date_forward,),])  # 合成因子是对date_forward下一期收益率的预测值
            df['tradedate'] = date_forward
            new_factor = pd.concat([new_factor, df])
        df = pd.DataFrame()
        df[f'{asset_type}code'] = asset_codes
        df['合成因子'] = np.mean(new_factor['合成因子'])
        df['tradedate'] = regress_dates[0]
        new_factor = pd.concat([new_factor, df])
        fac_data = new_factor.set_index(['tradedate', f'{asset_type}code'])
    elif method == 'pca':
        assert 'all_asset_quote' in kwargs, '请传入资产行情（all_asset_quote）变量，用于回归'
        regress_dates = factor_data_df.tradedate.unique()
        asset_codes = factor_data_df[f'{asset_type}code'].unique()
        stock_price_indexed = kwargs['all_asset_quote'].set_index(['tradedate', f'{asset_type}code'])
        stock_returns = stock_price_indexed.groupby(level=1).pct_change()
        fac_data_filled = fac_data.groupby(level=1).fillna(method='ffill')
        fac_data_filled = fac_data_filled.fillna(value=fac_data_filled.mean())
        new_factor = pd.DataFrame()

        # 普通PCA需要8个主成分才能解释90%以上的variance，而总共只有10个有效因子用于测试，可见常规PCA降维效果并不好
        # [0.27149649, 0.44775771, 0.56486784, 0.66704937, 0.7645456, 0.83978985, 0.89776431, 0.94574846]
        pca = decomposition.PCA(n_components=8)
        # pca.fit(fac_data_filled)
        # print(pca.explained_variance_ratio_)

        for i in range(len(regress_dates) - 1):
            df = pd.DataFrame()
            date = regress_dates[i]
            date_forward = regress_dates[i + 1]
            # centering处理
            X_modeling = fac_data_filled.loc[(date,),] - fac_data_filled.loc[(date,),].mean()
            X = pca.fit_transform(X_modeling)
            y = stock_returns.loc[(date_forward,),]

            model = LinearRegression().fit(X, y)
            df[f'{asset_type}code'] = asset_codes
            X_predicted = pca.fit_transform(
                fac_data_filled.loc[(date_forward,),] - fac_data_filled.loc[(date_forward,),].mean())
            df['合成因子'] = model.predict(X_predicted)  # 合成因子是对date_forward下一期收益率的预测值
            df['tradedate'] = date_forward
            new_factor = pd.concat([new_factor, df])
        df = pd.DataFrame()
        df[f'{asset_type}code'] = asset_codes
        df['合成因子'] = np.mean(new_factor['合成因子'])
        df['tradedate'] = regress_dates[0]
        new_factor = pd.concat([new_factor, df])
        fac_data = new_factor.set_index(['tradedate', f'{asset_type}code'])
    if method == 'external_weight':
        output_weight_df = weight_df.fillna(0)
    else:
        output_weight_df = None
    return fac_data.reset_index(), output_weight_df


if __name__ == '__main__':
    pass
